//===--- Job.cpp - Command to Execute -------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#include "polarphp/basic/StlExtras.h"
#include "polarphp/driver/DriverIncrementalRanges.h"
#include "polarphp/driver/Job.h"
#include "polarphp/driver/PrettyStackTrace.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Option/Arg.h"
#include "llvm/Support/Compiler.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/Program.h"
#include "llvm/Support/raw_ostream.h"

using namespace polar;
using namespace polar::driver;

StringRef CommandOutput::getOutputForInputAndType(StringRef PrimaryInputFile,
                                                  filetypes::FileTypeId Type) const {
   if (Type == filetypes::TY_Nothing)
      return StringRef();
   auto const *M = DerivedOutputMap.getOutputMapForInput(PrimaryInputFile);
   if (!M)
      return StringRef();
   auto const Out = M->find(Type);
   if (Out == M->end())
      return StringRef();
   assert(!Out->second.empty());
   return StringRef(Out->second);
}

struct CommandOutputInvariantChecker {
   CommandOutput const &Out;
   CommandOutputInvariantChecker(CommandOutput const &CO) : Out(CO) {
#ifndef NDEBUG
      Out.checkInvariants();
#endif
   }
   ~CommandOutputInvariantChecker() {
#ifndef NDEBUG
      Out.checkInvariants();
#endif
   }
};

void CommandOutput::ensureEntry(StringRef PrimaryInputFile,
                                filetypes::FileTypeId Type,
                                StringRef OutputFile,
                                bool Overwrite) {
   assert(!PrimaryInputFile.empty());
   assert(!OutputFile.empty());
   assert(Type != filetypes::TY_Nothing);
   auto &M = DerivedOutputMap.getOrCreateOutputMapForInput(PrimaryInputFile);
   if (Overwrite) {
      M[Type] = OutputFile;
   } else {
      auto res = M.insert(std::make_pair(Type, OutputFile));
      if (res.second) {
         // New entry, no need to compare.
      } else {
         // Existing entry, check equality with request.
         assert(res.first->getSecond() == OutputFile);
      }
   }
}

void CommandOutput::checkInvariants() const {
   filetypes::for_all_types([&](filetypes::FileTypeId Type) {
      size_t numOutputsOfType = 0;
      for (auto const &I : Inputs) {
         // FIXME: At the moment, empty primary input names correspond to
         // corner cases in the driver where it is doing TY_Nothing work
         // and isn't even given a primary input; but at some point we
         // ought to enable storing derived OFM entries under the empty
         // name in general, for "whole build" additional outputs. They
         // are presently (arbitrarily and wrongly) stored in entries
         // associated with the first primary input of the CommandOutput
         // that they were derived from.
         assert(PrimaryOutputType == filetypes::TY_Nothing || !I.Primary.empty());
         auto const *M = DerivedOutputMap.getOutputMapForInput(I.Primary);
         if (!M)
            continue;
         auto const Out = M->find(Type);
         if (Out == M->end())
            continue;
         assert(!Out->second.empty());
         ++numOutputsOfType;
      }
      assert(numOutputsOfType == 0 ||
             numOutputsOfType == 1 ||
             numOutputsOfType == Inputs.size());
   });
   assert(AdditionalOutputTypes.count(PrimaryOutputType) == 0);
}

bool CommandOutput::hasSameAdditionalOutputTypes(
   CommandOutput const &other) const {
   bool sameAdditionalOutputTypes = true;
   filetypes::for_all_types([&](filetypes::FileTypeId Type) {
      bool a = AdditionalOutputTypes.count(Type) == 0;
      bool b = other.AdditionalOutputTypes.count(Type) == 0;
      if (a != b)
         sameAdditionalOutputTypes = false;
   });
   return sameAdditionalOutputTypes;
}

void CommandOutput::addOutputs(CommandOutput const &other) {
   CommandOutputInvariantChecker Check(*this);
   assert(PrimaryOutputType == other.PrimaryOutputType);
   assert(&DerivedOutputMap == &other.DerivedOutputMap);
   Inputs.append(other.Inputs.begin(),
                 other.Inputs.end());
   // Should only be called with an empty AdditionalOutputTypes
   // or one populated with the same types as other.
   if (AdditionalOutputTypes.empty()) {
      AdditionalOutputTypes = other.AdditionalOutputTypes;
   } else {
      assert(hasSameAdditionalOutputTypes(other));
   }
}

CommandOutput::CommandOutput(filetypes::FileTypeId PrimaryOutputType,
                             OutputFileMap &Derived)
   : PrimaryOutputType(PrimaryOutputType), DerivedOutputMap(Derived) {
   CommandOutputInvariantChecker Check(*this);
}

filetypes::FileTypeId CommandOutput::getPrimaryOutputType() const {
   return PrimaryOutputType;
}

void CommandOutput::addPrimaryOutput(CommandInputPair Input,
                                     StringRef PrimaryOutputFile) {
   PrettyStackTraceDriverCommandOutputAddition CrashInfo(
      "primary", this, Input.Primary, PrimaryOutputType, PrimaryOutputFile);
   if (PrimaryOutputType == filetypes::TY_Nothing) {
      // For TY_Nothing, we accumulate the inputs but do not add any outputs.
      // The invariant holds on either side of this action because all primary
      // outputs for this command will be absent (so the length == 0 case in the
      // invariant holds).
      CommandOutputInvariantChecker Check(*this);
      Inputs.push_back(Input);
      return;
   }
   // The invariant holds in the non-TY_Nothing case before an input is added and
   // _after the corresponding output is added_, but not inbetween. Don't try to
   // merge these two cases, they're different.
   CommandOutputInvariantChecker Check(*this);
   Inputs.push_back(Input);
   assert(!PrimaryOutputFile.empty());
   assert(AdditionalOutputTypes.count(PrimaryOutputType) == 0);
   ensureEntry(Input.Primary, PrimaryOutputType, PrimaryOutputFile, false);
}

StringRef CommandOutput::getPrimaryOutputFilename() const {
   // FIXME: ideally this shouldn't exist, or should at least assert size() == 1,
   // and callers should handle cases with multiple primaries explicitly.
   assert(Inputs.size() >= 1);
   return getOutputForInputAndType(Inputs[0].Primary, PrimaryOutputType);
}

SmallVector<StringRef, 16> CommandOutput::getPrimaryOutputFilenames() const {
   SmallVector<StringRef, 16> V;
   size_t NonEmpty = 0;
   for (auto const &I : Inputs) {
      auto Out = getOutputForInputAndType(I.Primary, PrimaryOutputType);
      V.push_back(Out);
      if (!Out.empty())
         ++NonEmpty;
      assert(!Out.empty() || PrimaryOutputType == filetypes::TY_Nothing);
   }
   assert(NonEmpty == 0 || NonEmpty == Inputs.size());
   return V;
}

void CommandOutput::setAdditionalOutputForType(filetypes::FileTypeId Type,
                                               StringRef OutputFilename) {
   PrettyStackTraceDriverCommandOutputAddition CrashInfo(
      "additional", this, Inputs[0].Primary, Type, OutputFilename);
   CommandOutputInvariantChecker Check(*this);
   assert(Inputs.size() >= 1);
   assert(!OutputFilename.empty());
   assert(Type != filetypes::TY_Nothing);

   // If we're given an "additional" output with the same type as the primary,
   // and we've not yet had such an additional type added, we treat it as a
   // request to overwrite the primary choice (which happens early and is
   // sometimes just inferred) with a refined value (eg. -emit-module-path).
   bool Overwrite = Type == PrimaryOutputType;
   if (Overwrite) {
      assert(AdditionalOutputTypes.count(Type) == 0);
   } else {
      AdditionalOutputTypes.insert(Type);
   }
   ensureEntry(Inputs[0].Primary, Type, OutputFilename, Overwrite);
}

StringRef CommandOutput::getAdditionalOutputForType(filetypes::FileTypeId Type) const {
   if (AdditionalOutputTypes.count(Type) == 0)
      return StringRef();
   assert(Inputs.size() >= 1);
   // FIXME: ideally this shouldn't associate the additional output with the
   // first primary, but with a specific primary (and/or possibly the primary "",
   // for build-wide outputs) specified by the caller.
   assert(Inputs.size() >= 1);
   return getOutputForInputAndType(Inputs[0].Primary, Type);
}

SmallVector<StringRef, 16>
CommandOutput::getAdditionalOutputsForType(filetypes::FileTypeId Type) const {
   SmallVector<StringRef, 16> V;
   if (AdditionalOutputTypes.count(Type) != 0) {
      for (auto const &I : Inputs) {
         auto Out = getOutputForInputAndType(I.Primary, Type);
         // FIXME: In theory this should always be non-empty -- and V.size() would
         // always be either 0 or N like with primary outputs -- but in practice
         // WMO currently associates additional outputs with the _first primary_ in
         // a multi-primary job, which means that the 2nd..Nth primaries will have
         // an empty result from getOutputForInputAndType, and V.size() will be 1.
         if (!Out.empty())
            V.push_back(Out);
      }
   }
   assert(V.empty() || V.size() == 1 || V.size() == Inputs.size());
   return V;
}

StringRef CommandOutput::getAnyOutputForType(filetypes::FileTypeId Type) const {
   if (PrimaryOutputType == Type)
      return getPrimaryOutputFilename();
   return getAdditionalOutputForType(Type);
}

const OutputFileMap &CommandOutput::getDerivedOutputMap() const {
   return DerivedOutputMap;
}

StringRef CommandOutput::getBaseInput(size_t Index) const {
   assert(Index < Inputs.size());
   return Inputs[Index].Base;
}

static void escapeAndPrintString(llvm::raw_ostream &os, StringRef Str) {
   if (Str.empty()) {
      // Special-case the empty string.
      os << "\"\"";
      return;
   }

   bool NeedsEscape = Str.find_first_of(" \"\\$") != StringRef::npos;

   if (!NeedsEscape) {
      // This string doesn't have anything we need to escape, so print it directly
      os << Str;
      return;
   }

   // Quote and escape. This isn't really complete, but is good enough, and
   // matches how Clang's Command handles escaping arguments.
   os << '"';
   for (const char c : Str) {
      switch (c) {
         case '"':
         case '\\':
         case '$':
            // These characters need to be escaped.
            os << '\\';
            // Fall-through to the default case, since we still need to print the
            // character.
            LLVM_FALLTHROUGH;
         default:
            os << c;
      }
   }
   os << '"';
}

void
CommandOutput::print(raw_ostream &out) const {
   out
      << "{\n"
      << "    PrimaryOutputType = " << filetypes::get_type_name(PrimaryOutputType)
      << ";\n"
      << "    Inputs = [\n";
   interleave(Inputs,
              [&](CommandInputPair const &P) {
                 out << "        CommandInputPair {\n"
                     << "            Base = ";
                 escapeAndPrintString(out, P.Base);
                 out << ", \n"
                     << "            Primary = ";
                 escapeAndPrintString(out, P.Primary);
                 out << "\n        }";
              },
              [&] { out << ",\n"; });
   out << "];\n"
       << "    DerivedOutputFileMap = {\n";
   DerivedOutputMap.dump(out, true);
   out << "\n    };\n}";
}

void
CommandOutput::dump() const {
   print(llvm::errs());
   llvm::errs() << '\n';
}

void CommandOutput::writeOutputFileMap(llvm::raw_ostream &out) const {
   SmallVector<StringRef, 4> inputs;
   for (const CommandInputPair IP : Inputs) {
      assert(IP.Base == IP.Primary && !IP.Base.empty() &&
             "output file maps won't work if these differ");
      inputs.push_back(IP.Primary);
   }
   getDerivedOutputMap().write(out, inputs);
}

Job::~Job() = default;

void Job::printArguments(raw_ostream &os,
                         const llvm::opt::ArgStringList &Args) {
   interleave(Args,
              [&](const char *Arg) { escapeAndPrintString(os, Arg); },
              [&] { os << ' '; });
}

void Job::dump() const {
   printCommandLineAndEnvironment(llvm::errs());
}

ArrayRef<const char *> Job::getArgumentsForTaskExecution() const {
   if (hasResponseFile()) {
      writeArgsToResponseFile();
      return getResponseFileArg();
   } else {
      return getArguments();
   }
}

void Job::printCommandLineAndEnvironment(raw_ostream &Stream,
                                         StringRef Terminator) const {
   printCommandLine(Stream, /*Terminator=*/"");
   if (!ExtraEnvironment.empty()) {
      Stream << "  #";
      for (auto &pair : ExtraEnvironment) {
         Stream << " " << pair.first << "=" << pair.second;
      }
   }
   Stream << "\n";
}

void Job::printCommandLine(raw_ostream &os, StringRef Terminator) const {
   escapeAndPrintString(os, Executable);
   os << ' ';
   if (hasResponseFile()) {
      printArguments(os, {ResponseFile->argString});
      os << " # ";
   }
   printArguments(os, Arguments);

   os << Terminator;
}

void Job::printSummary(raw_ostream &os) const {
   // Deciding how to describe our inputs is a bit subtle; if we are a Job built
   // from a JobAction that itself has InputActions sources, then we collect
   // those up. Otherwise it's more correct to talk about our inputs as the
   // outputs of our input-jobs.
   SmallVector<StringRef, 4> Inputs;
   SmallVector<StringRef, 4> Outputs = getOutput().getPrimaryOutputFilenames();

   for (const Action *A : getSource().getInputs())
      if (const auto *IA = dyn_cast<InputAction>(A))
         Inputs.push_back(IA->getInputArg().getValue());

   for (const Job *J : getInputs())
      for (StringRef f : J->getOutput().getPrimaryOutputFilenames())
         Inputs.push_back(f);

   size_t limit = 3;
   size_t actual_in = Inputs.size();
   size_t actual_out = Outputs.size();
   if (actual_in > limit) {
      Inputs.erase(Inputs.begin() + limit, Inputs.end());
   }
   if (actual_out > limit) {
      Outputs.erase(Outputs.begin() + limit, Outputs.end());
   }

   os << "{" << getSource().getClassName() << ": ";
   interleave(Outputs,
              [&](const std::string &Arg) {
                 os << llvm::sys::path::filename(Arg);
              },
              [&] { os << ' '; });
   if (actual_out > limit) {
      os << " ... " << (actual_out-limit) << " more";
   }
   os << " <= ";
   interleave(Inputs,
              [&](const std::string &Arg) {
                 os << llvm::sys::path::filename(Arg);
              },
              [&] { os << ' '; });
   if (actual_in > limit) {
      os << " ... " << (actual_in-limit) << " more";
   }
   os << "}";
}


bool Job::writeArgsToResponseFile() const {
   assert(hasResponseFile());
   std::error_code EC;
   llvm::raw_fd_ostream OS(ResponseFile->path, EC, llvm::sys::fs::F_None);
   if (EC) {
      return true;
   }
   for (const char *arg : Arguments) {
      escapeAndPrintString(OS, arg);
      OS << " ";
   }
   OS.flush();
   return false;
}

StringRef Job::getFirstSwiftPrimaryInput() const {
   const JobAction &source = getSource();
   if (!isa<CompileJobAction>(source))
      return StringRef();
   const auto *firstInput = source.getInputs().front();
   if (auto *inputInput = dyn_cast<InputAction>(firstInput))
      return inputInput->getInputArg().getValue();
   return StringRef();
}

BatchJob::BatchJob(const JobAction &Source,
                   SmallVectorImpl<const Job *> &&Inputs,
                   std::unique_ptr<CommandOutput> Output,
                   const char *Executable, llvm::opt::ArgStringList Arguments,
                   EnvironmentVector ExtraEnvironment,
                   std::vector<FilelistInfo> Infos,
                   ArrayRef<const Job *> Combined, int64_t &NextQuasiPID,
                   Optional<ResponseFileInfo> ResponseFile)
   : Job(Source, std::move(Inputs), std::move(Output), Executable, Arguments,
         ExtraEnvironment, Infos, ResponseFile),
     CombinedJobs(Combined.begin(), Combined.end()),
     QuasiPIDBase(NextQuasiPID) {

   assert(QuasiPIDBase < 0);
   NextQuasiPID -= CombinedJobs.size();
   assert(NextQuasiPID < 0);
}
